#pragma once
#include<iostream>
#include<assert.h>

using std::cout;
using std::endl;
using std::pair;

template<class K,class V>
struct AVLTreeNode
{
    pair<K,V> _kv;
    AVLTreeNode<K,V>* _left;
    AVLTreeNode<K,V>* _right;
    AVLTreeNode<K,V>* _parent;

    int _bf;//balance factor

    AVLTreeNode(const pair<K,V>& kv)
    :_kv(kv)
    ,_left(nullptr)
    ,_right(nullptr)
    ,_parent(nullptr)
    {}
};

template<class K,class V>
class AVLTree
{
    typedef AVLTreeNode<K,V> Node;
public:
    AVLTree():_root(nullptr){}

    bool Insert(const pair<K,V>& kv)
    {
        if(_root==nullptr)
        {
            _root=new Node(kv);
            return true;
        }
        //1.找位置，同时记录cur的父节点
        Node* cur=_root,*parent=nullptr;
        while(cur)
        {
            if(kv.first>cur->_kv.first)
                parent=cur,cur=cur->_right;
            else if(kv.first<cur->_kv.first)
                parent=cur,cur=cur->_left;
            else return false;
        }
        //2.插入结点,此时parent必定不为空
        cur=new Node(kv);
        if(kv.first>parent->_kv.first)
            parent->_right=cur,cur->_parent=parent;
        else 
            parent->_left=cur,cur->_parent=parent;

        //向上更新_bf
        while(parent)
        {
            if(cur==parent->_right)
                parent->_bf++;
            else 
                parent->_bf--;

            //判断父节点的_bf
            if(parent->_bf==0)
                break; //跳出循环
            else if(parent->_bf==1 || parent->_bf==-1)
                cur=parent,parent=parent->_parent;   //继续向上更新
            else if(parent->_bf==2 || parent->_bf==-2)
            {
                //旋转，调整_bf
                if(parent->_bf==2 && cur==parent->_right && cur->_bf==1)
                    RotateL(parent);//左单旋
                else if(parent->_bf==-2 && cur==parent->_left && cur->_bf==-1)
                    RotateR(parent);//右单旋
                else if(parent->_bf==2 && cur==parent->_right && cur->_bf==-1)
                    RotateRL(parent);//右左双旋
                else if(parent->_bf==-2 && cur==parent->_left && cur->_bf==1)
                    RotateLR(parent);//左右双旋
                break;
            }
            else assert(false);
        }
        return true;
    }

    void Inorder()
    {
        _Inorder(_root);
        cout<<endl;
    }

    bool Is_Balance()
    {
        return _Is_Balance(_root);
    }

    int Height()
    {
        return _Height(_root);
    }

private:
    void RotateL(Node* parent)
    {
        Node* grandparent=parent->_parent;

        Node* cur=parent->_right;
        Node* curleft=cur->_left;

        parent->_right=curleft;
        parent->_parent=cur;

        if(curleft)
            curleft->_parent=parent;
        
        cur->_left=parent;
        cur->_parent=grandparent;

        if(grandparent==nullptr)
            _root=cur;
        else 
        {
            if(parent==grandparent->_left)
                grandparent->_left=cur;
            else 
                grandparent->_right=cur;
        }

        cur->_bf=parent->_bf=0;
    }

    void RotateR(Node* parent)
    {
        Node* grandparent=parent->_parent;

        Node* cur=parent->_left;
        Node* curright=cur->_right;

        parent->_left=curright;
        parent->_parent=cur;

        if(curright)
            curright->_parent=parent;
        
        cur->_right=parent;
        cur->_parent=grandparent;

        if(grandparent==nullptr)
            _root=cur;
        else 
        {
            if(parent==grandparent->_left)
                grandparent->_left=cur;
            else 
                grandparent->_right=cur;
        }

        cur->_bf=parent->_bf=0;
    }

    void RotateRL(Node* parent)
    {
        Node* cur=parent->_right;
        Node* curleft=cur->_left;//curleft必定存在
        int bf=curleft->_bf;

        RotateR(cur);
        RotateL(parent);

        //经过右、左单旋后，cur、parent、curleft的_bf都被置为0
        //根据旋转前记录的curleft的bf判断cur和parent的_bf
        if(bf==0)
            cur->_bf=parent->_bf=0;
        else if(bf==1)
            cur->_bf=0,parent->_bf=-1;
        else if(bf==-1)
            cur->_bf=1,parent->_bf=0;
    }

    void RotateLR(Node* parent)
    {
        Node* cur=parent->_left;
        Node* curright=cur->_right;//curleft必定存在
        int bf=curright->_bf;

        RotateL(cur);
        RotateR(parent);

        //经过右、左单旋后，cur、parent、curright的_bf都被置为0
        //根据旋转前记录的curright的bf判断cur和parent的_bf
        if(bf==0)
            cur->_bf=parent->_bf=0;
        else if(bf==1)
            cur->_bf=-1,parent->_bf=0;
        else if(bf==-1)
            cur->_bf=0,parent->_bf=1;
    }

    void _Inorder(Node* root)
    {
        if(root==nullptr)
            return;
        _Inorder(root->_left);
        cout<<root->_kv.first<<" ";
        _Inorder(root->_right);
    }

    int _Height(Node* root)
    {
        if(root==nullptr)
            return 0;

        int left_height=_Height(root->_left);
        int right_height=_Height(root->_right);

        return left_height>right_height?left_height+1:right_height+1;
    }

    bool _Is_Balance(Node* root)
    {
        if(root==nullptr)
            return true;
        //判断一棵AVL树是否平衡：每个结点的abs(右子树高度-左子树高度)<=1
        int left_height=_Height(root->_left);
        int right_height=_Height(root->_right);

        if(abs(right_height-left_height)>1 || right_height-left_height!=root->_bf)
        {
            cout<<root->_kv.first<<"->"<<root->_bf<<"...error!"<<endl;
            return false;
        }

        return _Is_Balance(root->_left) && _Is_Balance(root->_right);
    }

private:
    Node* _root;
};